package com.moon.cloud.common.utils;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 树形工具类
 */
public class TreeUtils {

    /**
     * 构建树形结构（顺序执行）
     */
    public static <T> List<T> buildTree(List<T> list,
                                        Function<T, Object> idGetter,
                                        Function<T, Object> parentIdGetter,
                                        Function<T, List<T>> childrenGetter,
                                        BiConsumer<T, List<T>> childrenSetter) {
        return buildTree(list, idGetter, parentIdGetter, childrenGetter, childrenSetter, false);
    }

    /**
     * 构建树形结构
     *
     * @param list           待构建的列表
     * @param idGetter       获取ID的函数
     * @param parentIdGetter 获取父ID的函数
     * @param childrenGetter 获取子节点列表的函数
     * @param childrenSetter 设置子节点列表的函数
     * @param isParallel     是否并行执行
     */
    public static <T> List<T> buildTree(List<T> list,
                                        Function<T, Object> idGetter,
                                        Function<T, Object> parentIdGetter,
                                        Function<T, List<T>> childrenGetter,
                                        BiConsumer<T, List<T>> childrenSetter,
                                        Boolean isParallel) {
        if (list == null || list.isEmpty()) {
            return Collections.emptyList();
        }

        // 构建 ID -> 节点映射
        Map<Object, T> nodeMap = isParallel
                ? list.parallelStream().collect(Collectors.toConcurrentMap(idGetter, Function.identity(), (existing, _) -> existing, ConcurrentHashMap::new))
                : list.stream().collect(Collectors.toMap(idGetter, Function.identity(), (existing, _) -> existing, LinkedHashMap::new));

        // 构建树形结构
        return nodeMap.values().stream()
                .peek(node -> {
                    Object parentId = parentIdGetter.apply(node);
                    if (!isRootNode(parentId) && nodeMap.containsKey(parentId)) {
                        T parent = nodeMap.get(parentId);
                        List<T> children = getOrCreateChildren(parent, childrenGetter, childrenSetter);
                        children.add(node);
                    }
                })
                .filter(node -> isRootNode(parentIdGetter.apply(node)))
                .collect(Collectors.toList());
    }

    /**
     * 判断是否为根节点（parentId 为 null 或 0）
     */
    private static boolean isRootNode(Object parentId) {
        if (parentId == null) {
            return true;
        }
        // 支持 Long、Integer、String 类型的 0
        if (parentId instanceof Number) {
            return ((Number) parentId).longValue() == 0L;
        }
        return false;
    }

    /**
     * 获取或创建子节点列表
     */
    private static <T> List<T> getOrCreateChildren(T parent,
                                                   Function<T, List<T>> childrenGetter,
                                                   BiConsumer<T, List<T>> childrenSetter) {
        List<T> children = childrenGetter.apply(parent);
        if (children == null) {
            children = new ArrayList<>();
            childrenSetter.accept(parent, children);
        }
        return children;
    }

}
